# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question1/models.py
import torch.nn

import torch_plus


# noinspection SpellCheckingInspection
class LorentzDispersionModel(torch.nn.Module):
    # 此默认值是由 main_guess.py 给出的
    def __init__(self,
                 epsilon_infinity: float = 3.83,
                 omega_p: float = 15.94,
                 omega_alpha: float = 5.4,
                 omega_c: float = 2.3):
        super().__init__()
        self.epsilon_infinity = torch.nn.Parameter(
            torch_plus.as_tensor(epsilon_infinity))
        self.omega_p = torch.nn.Parameter(
            torch_plus.as_tensor(omega_p))
        self.omega_alpha = torch.nn.Parameter(
            torch_plus.as_tensor(omega_alpha))
        self.omega_c = torch.nn.Parameter(
            torch_plus.as_tensor(omega_c))
        self.__two = torch_plus.as_tensor(2)
        self.__j = torch.tensor(1j)

    def forward(self, omega):
        denominator = (self.omega_alpha ** self.__two
                       - self.__j * self.omega_c * omega
                       - omega ** self.__two)
        epsilon = self.epsilon_infinity + self.omega_p ** self.__two / denominator
        return torch.stack([torch.real(epsilon), torch.imag(epsilon)], dim=-1)
